import numpy
import scipy
from math import sqrt
from kuai.sim.simtool import *
from kuai.sim.engine import *

class LeapFrogControl(Control):
    def run(self, engine, mol, ff, energy, i, step_length):
        half_step = step_length / 2.0
        mol.coords += mol.speeds * half_step
        mol.forces = -energy.grediant(mol, ff)
        acceleration = mol.forces * mol.rweights         # acc = Force / Mass 
        mol.speeds += acceleration * step_length
        mol.coords += mol.speeds * half_step

class MoleculeDynamicEngine(SimulationEngine):
    def __init__(self, mol, ff):
        SimulatinEngine.__init__(self)
        MoleculeDynamicEngine.__prepare_md(mol, ff)
        self.use(LeapFrogControl())
    
    @staticmethod
    def __prepare_md(mol, ff):
        mol.weights = collect_coords(mol, ff)
        mol.rweights = 1.0 / mol.weights    # reversed weights
        ff.k = collect_k(ff)

        natoms = len(mol.atoms())
        mol.coords = collect_coords(mol)
        mol.speeds = numpy.zeros(shape=(natoms*3, 1), dtype=REAL_TYPE)   # The speeds
        mol.forces = numpy.zeros(shape=(natoms*3, 1), dtype=REAL_TYPE)   # The forces
        
    def run(self, mol, ff, energy, n, step_length = 1, *args, **kwargs):
        SimulatinEngine.run(self, mol, ff, energy, n, step_length, *args, **kwargs)


class TemperatureControl(Control):
    def __init__(self, t):
        self.temperature = t    # Target Temperature to Control
    
    @staticmethod
    def get_temperature(self, mol):
        v = numpy.dot(mol.speeds, mol.speeds * mol.weights) / 2
        return v
    
    @staticmethod
    def init_temperature(self, mol, t):
        pass
    
    def run(self, engine, mol, ff, energy, i, step_length):
        pass

class BerendsenTemperatureControl(TemperatureControl):
    def __init__(self, t, tau = 100):
        TemperatureControl.__init__(self, t)
        self.__rtau = 1.0 / tau
        
    def run(self, engine, mol, ff, energy, i, step_length):
        origin = self.get_temperature(mol)
        factor = sqrt( 1 + step_length * _compressibility * (self.temperature/origin - 1 ));
        mol.speed *= factor
        mol.temperature = origin * factor * factor
        
class AndersonTemperatureControl1(TemperatureControl):
    def __init__(self, t):
        TemperatureControl.__init__(self, t)
        
    def run(self, engine, mol, ff, energy, i, *argv, **kwarg):
        self.init_temperature(mol, self.temperature)
        
class BerendsenTemperatureControl(TemperatureControl):
    def __init__(self, t, tau = 100):
        TemperatureControl.__init__(self, t)
        self.__rtau = 1.0 / tau
        
    def run(self, engine, mol, ff, energy, i, step_length):
        origin = self.get_temperature(mol)
        factor = 1 + step_length * self.__rtau * (self.temperature/origin - 1 );
        mol.speed *= sqrt(factor)
        mol.temperature = origin * factor

class PressureControl(Control):
    def __init__(self, p):
        self.pressure = p   # Target Pressure to Control
        
class BerendsenPressureControl(PressureControl):
    def __init__(self, p, tau=100):
        PressureControl.__init__(self, p)
        self.__rtau = 1.0 / tau
        
    def run(self, engine, mol, ff, energy, i, step_length):
        origin = self.get_pressure(mol)
        factor = 1 + step_length * self.__rtau * (self.pressure/origin - 1 );
        mol.coords *= factor ** (1.0 / 3)
        mol.pressure = origin * factor

